package websocket

import (
	"sync"

	"github.com/gorilla/websocket"
)

// Client represents a WebSocket client connection
type Client struct {
	UserID int64
	Conn   *websocket.Conn
	Send   chan []byte
}

// Message represents a WebSocket message
type Message struct {
	UserID  int64
	Payload []byte
}

// Hub maintains the set of active clients and broadcasts messages
type Hub struct {
	clients    map[int64]map[*Client]bool // Map of userID to their connections
	register   chan *Client
	unregister chan *Client
	broadcast  chan *Message
	mu         sync.RWMutex
}

// GlobalHub is the global WebSocket hub instance
var GlobalHub = NewHub()

// NewHub creates a new hub instance
func NewHub() *Hub {
	return &Hub{
		clients:    make(map[int64]map[*Client]bool),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		broadcast:  make(chan *Message),
	}
}

// Run starts the hub's message processing loop
func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			if _, exists := h.clients[client.UserID]; !exists {
				h.clients[client.UserID] = make(map[*Client]bool)
			}
			h.clients[client.UserID][client] = true
			h.mu.Unlock()
		case client := <-h.unregister:
			h.mu.Lock()
			if conns, ok := h.clients[client.UserID]; ok {
				if _, ok := conns[client]; ok {
					delete(conns, client)
					close(client.Send)
					if len(conns) == 0 {
						delete(h.clients, client.UserID)
					}
				}
			}
			h.mu.Unlock()
		case msg := <-h.broadcast:
			h.mu.RLock()
			if msg.UserID == 0 {
				// Broadcast to all clients
				for _, conns := range h.clients {
					for client := range conns {
						select {
						case client.Send <- msg.Payload:
						default:
							close(client.Send)
							delete(conns, client)
						}
					}
				}
			} else {
				// Send to specific user's clients
				if conns, ok := h.clients[msg.UserID]; ok {
					for client := range conns {
						select {
						case client.Send <- msg.Payload:
						default:
							close(client.Send)
							delete(conns, client)
						}
					}
				}
			}
			h.mu.RUnlock()
		}
	}
}

// BroadcastToUser 推送消息给指定用户
func (h *Hub) BroadcastToUser(userID int64, payload []byte) {
	h.broadcast <- &Message{
		UserID:  userID,
		Payload: payload,
	}
}
